use std::ops::Deref;

use humanize_rs::bytes::{Bytes, Unit};
use humansize::{format_size, DECIMAL};
use ratatui::widgets::Row;
use serde::Deserialize;
use string_enum::StringEnum;
use urlparse::GetQuery;

use crate::tui::ui::TableData;

#[derive(Deserialize, StringEnum, Clone)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ModelType {
    /// `Checkpoint`
    Checkpoint,
    /// `Lora`
    Lora,
    /// `Lycois`
    Lycois,
    /// `Embedding`
    Embedding,
    /// `HyperNetwork`
    HyperNetwork,
    /// `ControlNet`
    ControlNet,
    /// `VAE`
    Vae,
    /// `ESRGAN`
    Esrgan,
}

#[derive(Deserialize, StringEnum, Clone)]
pub(crate) enum BaseModelType {
    /// `sd1.5`
    #[serde(rename = "sd1.5")]
    SD15,
    /// `sd2.1`
    #[serde(rename = "sd2.1")]
    SD21,
    /// `sdxl`
    #[serde(rename = "sdxl")]
    SDXL,
}

#[derive(Deserialize, StringEnum, Clone)]
#[serde(rename_all = "snake_case")]
pub(crate) enum SourceSite {
    /// `Civitai`
    Civitai,
    /// `LibLibAI`
    Liblibai,
    /// `XGY`
    Xgy,
}

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct Model {
    pub name: Option<String>,
    #[serde(rename = "type")]
    pub model_type: ModelType,
    pub url: String,
    pub cover: Option<String>,
    pub extra: Option<String>,
    pub source: SourceSite,
    #[serde(with = "serde_humanize_rs")]
    pub size: Option<usize>,
    pub active_words: Option<Vec<String>>,
    pub base_model: BaseModelType,
}

impl Model {
    pub fn get_model_name(&self) -> String {
        if let Some(name) = &self.name {
            return name.clone();
        }
        let url = urlparse::urlparse(self.url.as_str());
        let path = url.path.clone();
        match self.source {
            SourceSite::Civitai | SourceSite::Xgy => path.split("/").last().unwrap().to_string(),
            SourceSite::Liblibai => match &url.get_parsed_query() {
                Some(q) => q.get_first_from_str("attname").unwrap().to_string(),
                None => path.split("/").last().unwrap().to_string(),
            },
        }
    }

    pub fn get_extra_name(&self) -> Option<String> {
        match &self.extra {
            Some(extra) => {
                let url = urlparse::urlparse(extra.as_str());
                let path = url.path.clone();
                Some(path.split("/").last().unwrap().to_string())
            }
            None => None,
        }
    }
}

impl<'a> TableData<'a> for Model {
    type Data = Self;
    type TableData = Row<'a>;

    fn data(&self) -> &Self::Data {
        &self
    }

    fn to_table_data(&self) -> Self::TableData {
        let active_words = self.active_words.clone();
        Row::new(vec![
            self.get_model_name(),
            self.model_type.to_string(),
            self.base_model.to_string(),
            active_words.unwrap_or(vec![]).join(","),
            format_size(self.size.unwrap_or(0), DECIMAL),
            self.source.to_string(),
        ])
    }
}

#[derive(Deserialize, Debug)]
pub(crate) struct ModelList {
    models: Vec<Model>,
}

impl Deref for ModelList {
    type Target = Vec<Model>;

    fn deref(&self) -> &Self::Target {
        return &self.models;
    }
}

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct ModelPaths {
    checkpoint: String,
    controlnet: String,
    embedding: String,
    esrgan: String,
    hypernetworks: String,
    lora: String,
    lycois: String,
    vae: String,
}

impl ModelPaths {
    pub fn get(&self, t: &ModelType) -> Option<&String> {
        match t {
            ModelType::Checkpoint => Some(&self.checkpoint),
            ModelType::ControlNet => Some(&self.controlnet),
            ModelType::Embedding => Some(&self.embedding),
            ModelType::Esrgan => Some(&self.esrgan),
            ModelType::HyperNetwork => Some(&self.hypernetworks),
            ModelType::Lora => Some(&self.lora),
            ModelType::Lycois => Some(&self.lycois),
            ModelType::Vae => Some(&self.vae),
        }
    }
}

#[derive(Deserialize, Debug, Clone)]
pub(crate) struct SDConfig {
    pub path: String,
    pub model_paths: ModelPaths,
    pub models: Vec<Model>,
}

impl<'a> TableData<'a> for SDConfig {
    type TableData = Vec<Row<'a>>;
    type Data = Vec<Model>;

    fn data(&self) -> &Self::Data {
        &(self.models)
    }

    fn to_table_data(&self) -> Self::TableData {
        self.models
            .iter()
            .map(|m| m.to_table_data())
            .collect::<Vec<Row>>()
    }
}
